import cv2
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import h5py as h5

BASE_DIR = Path(__file__).resolve().parent


def init_plot():
    fig = plt.figure(figsize=(16,10))
    fig.canvas.mpl_connect('key_press_event', on_press)
    fig.subplots_adjust(left=0.02, right=0.98, bottom=0.1, top=0.9)
    axes = []
    axes.append(fig.add_subplot(2, 3, 1))
    axes.append(fig.add_subplot(2, 3, 2))
    axes.append(fig.add_subplot(2, 3, 3))
    axes.append(fig.add_subplot(2, 3, 4))
    axes.append(fig.add_subplot(2, 3, 5))
    axes.append(fig.add_subplot(2, 3, 6))
    axes[0].set_title('left',fontsize=12, pad=13)
    axes[1].set_title('left_with_marker',fontsize=12, pad=13)
    axes[2].set_title('left_with_marker_sim',fontsize=12, pad=13)
    axes[3].set_title('right',fontsize=12, pad=13)
    axes[4].set_title('right_with_marker', fontsize=12, pad=13)
    axes[5].set_title('right_with_marker_sim', fontsize=12, pad=13)
    return fig, axes

def plot(axes, imgs):
    for i, img in enumerate(imgs):
        axes[i].imshow(img)
    plt.show(block=True)

def on_press(event):
    if event.key == "enter":
        plt.close()
    elif event.key == "escape":
        plt.close()
        exit()


if __name__ == '__main__':

    with h5.File( 'G:/data.h5', 'r') as f:
        for i in range(100):
            left_rgb = f['left_rgb'][i]
            left_rgb_with_marker = f['left_rgb_with_marker'][i]
            left_sim = f['left_rgb_with_marker_sim'][i]

            right_rgb = f['right_rgb'][i]
            right_rgb_with_marker = f['right_rgb_with_marker'][i]
            right_sim = f['right_rgb_with_marker_sim'][i]
            plot(
                init_plot()[1],
                [left_rgb, left_rgb_with_marker, left_sim, right_rgb, right_rgb_with_marker, right_sim]
            )